from audioop import cross, reverse
import numpy as np
import random
import copy

class Individual:

    def __init__(self, obj_index, inner_id, config):
        """_summary_
        Class Individual implement the single Individual (in one swarm for objective i) actions.

        Args:
            obj_index (int): the i-th objective for the i-th swarm mainly concerned.
            inner_id (int): the j-th Individual in the i-th swarm.
            config (dict): include the information about the decision space dimesion, objective space dimesion and so on.
        """

        self.obj_index = obj_index
        self.inner_id = inner_id
        self.config = config

        self.F1 = self.config['F1']
        self.F2 = self.config['F2']
        self.CR = self.config['CR']
        
        ## solution
        self.x = self.config['LowerBound'] + (self.config['UpperBound'] - self.config['LowerBound']) * np.random.rand(self.config['DsDim'])        
        ## fitness
        self.fitness = self.config['evaluationer'].eval(self.x)
        
        self.distance = 0.0  ## for crowding ditance sort

    def update_Individual(self, group_A, group_B):
        """_summary_

        Args:
            archive (Dict): The external archive for information sharing, using two elements, the 'size' and the 'set'. 
        """

        ## mutation        
        v = group_A[0] + self.F1 * (group_A[1] - group_A[2]) + self.F2 * (group_B[0] - group_B[1])
        # v = group_A[0] + self.F1 * (group_A[1] - group_A[2])
        v = np.clip(v, a_min=self.config['LowerBound'], a_max=self.config['UpperBound'])

        ## crossover
        cross_x = np.zeros_like(v)
        for j in range(self.config['DsDim']):
            rand_j = random.randint(0, self.config['DsDim'] - 1)
            rand_float = random.random()
            if rand_float <= self.CR or rand_j == j:
                cross_x[j] = v[j]
            else:
                cross_x[j] = self.x[j]
        cross_fitness = self.config['evaluationer'].eval(cross_x)
        
        ## selection
        if cross_fitness[self.obj_index] < self.fitness[self.obj_index]:
            self.x = cross_x
            self.fitness = cross_fitness

    def deepcopy(self, copy_object):
        """_summary_
        Let the copy_objectt(a Individual)'s information be copyed to 'this' Individual.

        Args:
            copy_objec (_type_): type: Individual
        """
        self.obj_index = copy_object.obj_index
        self.x = copy.deepcopy(copy_object.x)
        self.fitness = copy.deepcopy(copy_object.fitness)
    
    def els(self, ):

        d = random.randint(0, self.config['DsDim'] - 1)
        self.x[d] += (self.config['UpperBound'] - self.config['LowerBound']) * random.gauss(0., 1.)
        
        ## here, we following the offical code (C plus plus version), if the x[d] out of range, we randomly sample the x[d] into [lower, upper]
        if self.x[d] < self.config['LowerBound'] or self.x[d] > self.config['UpperBound']:
            self.x[d] = random.uniform(self.config['LowerBound'], self.config['UpperBound'])

        self.fitness = self.config['evaluationer'].eval(self.x)

    def reset_distance(self, ):
        self.distance = 0.0

class CMDE:

    def __init__(self, config):
        """_summary_
        The main body for CMPSO.
        the main process:
            1) maintain a archive.
            2) build a swarms, each swarm have a mainly objective and maintain some population in the swarm.
        Args:
            config (Dict): Include the information about the decision space dimesion, objective space dimesion and so on.
        """
        
        self.config = config
        
        self.DsDim = self.config['DsDim']
        self.OsDim = self.config['OsDim']
        
        ## Init the swarms, where each swarm mainly for one objective.
        self.swarms = []
        for obj_index in range(self.OsDim):
            #swarm for i-th objective
            obj_i_swarm = [Individual(obj_index=obj_index, inner_id=i, config=self.config) for i in range(self.config['NP'])]
            if len(obj_i_swarm) != self.config['NP']:
                raise ValueError

            self.swarms.append(obj_i_swarm)


        ## Init the archive
        self.archive = {
            'size': 0,  ## current size
            'upper_size': self.config['arc_size'], ## the upper size of the archive
            'set': [] ## to stroe those Individual
        }
        self.update_archive()


    def update_archive(self, ):
        
        def dominate(fitness_1, fitness_2):
            """
            return 0/-1/1, where 0 represent co-dominate, -1 represent fit_2 domainte fit_1, 1 represent fit 1 domainate fit_2.
            """
            better = (fitness_1 < fitness_2).sum()
            same = (fitness_1 == fitness_2).sum()
            worse = (fitness_1 > fitness_2).sum()

            if same == fitness_1.shape[0]:
                return 0
            elif better == 0:
                return -1
            elif worse == 0:
                return 1
            else:
                return 0

        def nondominatedSet(S_set):
            """_summary_
                Get non-dominated Set R and return.
            Args:
                S_set (_type_): _description_
            Retrun:
                R_set (List):
            """
            size = len(S_set)
            R_set = []

            for i in range(size):
                non = True
                individual_a = S_set[i]
                for j in range(size):
                    if i != j:
                        individual_b = S_set[j]
                        k = dominate(individual_a.fitness, individual_b.fitness)
                        if k < 0:
                            non = False
                            break
                if non:
                    R_set.append(individual_a)
            return R_set            

        def crowding_distance_assignment(R_set):
            """_summary_
                Give each Individuals a crowding distance, and return the "sorted(by distance)" set.
            Args:
                R_set (List): 
            """
            ## zero the distance for Individuals in R_set
            size = len(R_set)
            for Individual in R_set:
                Individual.reset_distance()
            for sort_obj_index in range(self.OsDim):
                # sort by the i-th objective
                R_set.sort(key=lambda Individual: Individual.fitness[sort_obj_index])
                R_set[0].distance = 1e6
                R_set[size - 1].distance = 1e6
                fitmin = R_set[0].fitness[sort_obj_index]
                fitmax = R_set[size - 1].fitness[sort_obj_index]

                for index in range(1, size - 1):
                    R_set[index].distance += (R_set[index + 1].fitness[sort_obj_index] - R_set[index - 1].fitness[sort_obj_index]) / (fitmax - fitmin)
            
            # sort by distance
            R_set.sort(key=lambda Individual: Individual.distance, reverse=True)

            return R_set

        S = [] 
        
        ## Add pBest of each Individual in each swarm into set s
        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                S.append(self.swarms[obj_index][i])


        ## Add all the solution in the old archive into the set S
        for i in range(self.archive['size']):
            S.append(self.archive['set'][i])
        

        ## ELS strategy
        for i in range(self.archive['size']):
            temp = Individual(obj_index=self.archive['set'][i].obj_index, inner_id=-1, config=self.config)
            temp.deepcopy(self.archive['set'][i])
            temp.els()
            S.append(temp)
        
        R = nondominatedSet(S)

        if len(R) < self.archive['upper_size']:
            self.archive['set'] = copy.deepcopy(R)
            
            if len(self.archive['set']) != len(R):
                raise ValueError
            
            self.archive['size'] = len(self.archive['set'])
        else:
            sorted_R = crowding_distance_assignment(R)
            self.archive['set'].clear()
            for i in range(self.archive['upper_size']):
                self.archive['set'].append(sorted_R[i])
            if len(self.archive['set']) != self.archive['upper_size']:
                raise ValueError
            self.archive['size'] = len(self.archive['set'])    

            del sorted_R
        
        del S 
        del R   
        
    def update_swarms(self, ):

        ## Following the paper "Multiple Populations for Multiple Objectives: A Coevolutionary Technique for Solving Multiobjective Optimization Problems" Figure3
        ## Select an achieve solution randomly for each Individual
        ## Update velocity and position using eq.6 and eq.7
        ## Eval
        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                r1 = r2 = r3 = 0
                while r1 == i or r2 == i or r3 == i or r2 == r1 or r3 == r1 or r3 == r2:
                    r1 = random.randint(0, self.config['NP'] - 1)
                    r2 = random.randint(0, self.config['NP'] - 1)
                    r3 = random.randint(0, self.config['NP'] - 1)
                group_A = []
                group_A.append(self.swarms[obj_index][r1].x)
                group_A.append(self.swarms[obj_index][r2].x)
                group_A.append(self.swarms[obj_index][r3].x)
                
                r4 = r5 = 0
                group_B = []
                if self.archive['size'] < 2:
                    while r4 == r5:
                        r4 = random.randint(0, self.config['NP'] - 1)
                        r5 = random.randint(0, self.config['NP'] - 1)
                    group_B.append(self.swarms[(obj_index + 1) % 2][r4].x)
                    group_B.append(self.swarms[(obj_index + 1) % 2][r5].x)
                else:
                    while r4 == r5:
                        r4 = random.randint(0, self.archive['size'] - 1)
                        r5 = random.randint(0, self.archive['size'] - 1)
                    group_B.append(self.archive['set'][r4].x)
                    group_B.append(self.archive['set'][r5].x)

                self.swarms[obj_index][i].update_Individual(group_A, group_B)
